import torch
import shutil
import os
from utils.oss_file import download_file


class Solver:
    def __init__(self, model, device):
        self.model = model
        self.device = device

    def forward(self, images):
        images = images.to(self.device)
        outputs = self.model(images)
        return outputs

    def tta(self, images, seg=True):
        """
        测试时数据增强，对待预测图像进行水平、垂直翻转再预测，将预测结果翻转回去之后和原图的预测结果相加求平均
        :param images: [b,c,h,w]
        :param seg: 分类还是分割，默认为分割
        :return:
        """
        images = images.to(self.device)  # 原图
        pred_origin = self.model(images)  # 原图预测

        preds = torch.zeros_like(pred_origin)

        # 水平翻转
        images_hflp = torch.flip(images, dims=[3])
        pred_hflip = self.model(images_hflp)
        # 垂直翻转
        images_vflip = torch.flip(images, dims=[2])
        pred_vflip = self.model(images_vflip)
        # 水平加垂直翻转
        # images_hvflip = torch.flip(images, dims=[3])
        # images_hvflip = torch.flip(images_hvflip, dims=[2])
        # pred_hvflip = self.model(images_hvflip)

        if seg:
            # 分割需要将预测结果翻转回去
            pred_hflip = torch.flip(pred_hflip, dims=[3])
            pred_vflip = torch.flip(pred_vflip, dims=[2])
            # pred_hvflip = torch.flip(pred_hvflip, dims=[2])
            # pred_hvflip = torch.flip(pred_hvflip, dims=[3])
        preds = preds + pred_origin + pred_hflip + pred_vflip  # + pred_hvflip
        # 求平均
        pred = preds / 3.0

        return pred

    def cal_loss(self, targets, predicts, criterion):
        targets = targets.to(self.device)
        return criterion(predicts, targets)

    def backword(self, optimizer, loss):
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    def save_checkpoint(self, save_path, state, is_best, logger):
        if is_best:
            logger.log('Saving Best Model...')
            save_best_path = save_path.replace('.pth', '_best.pth')
            state['logger_next_position'] = logger.next_position
            torch.save(state, save_best_path)
        torch.save(state, save_path)

    def load_checkpoint(self, load_path):
        if os.path.isfile(load_path):
            checkpoint = torch.load(load_path, map_location=self.device)
            self.model.load_state_dict(checkpoint['state_dict'])
            print(f'Successfully Loaded from {load_path}')
            # return self.model
        else:
            raise FileNotFoundError("Can not find weight file in {}".format(load_path))